import torch.nn as nn
from torch_geometric.graphgym.config import cfg
from torch_geometric.graphgym.models.layer import new_layer_config, MLP
from torch_geometric.graphgym.register import register_head


@register_head('inductive_node')
class GNNInductiveNodeHead(nn.Module):
    """
    GNN prediction head for inductive node prediction tasks.

    Args:
        dim_in (int): Input dimension
        dim_out (int): Output dimension. For binary prediction, dim_out=1.
    """

    def __init__(self, dim_in, dim_out):
        super(GNNInductiveNodeHead, self).__init__()
        # self.layer_post_mp: 定义后多层感知机（post-message-passing）的模型层，用于在 GNN 模型的消息传递结束后对数据进一步处理。
        # MLP(...): 使用 torch_geometric 提供的 MLP 模块创建一个多层感知机。
        # new_layer_config(...): 配置 MLP 的层次结构。
        # dim_in: 输入维度。
        # dim_out: 输出维度。
        # cfg.gnn.layers_post_mp: 指定 post_mp 层的数量。
        # has_act=False: 指定不使用激活函数。
        # has_bias=True: 指定层是否使用偏置项。
        # cfg=cfg: 传递全局配置对象。
        self.layer_post_mp = MLP(
            new_layer_config(dim_in, dim_out, cfg.gnn.layers_post_mp,
                             has_act=False, has_bias=True, cfg=cfg))

    # 这是一个辅助函数，用于从 batch 中提取特征（x）和标签（y）。假设 batch.x 是节点特征，batch.y 是标签。
    def _apply_index(self, batch):
        return batch.x, batch.y

    def forward(self, batch):
        batch = self.layer_post_mp(batch)
        pred, label = self._apply_index(batch)
        return pred, label
